mirror of
https://github.com/zama-ai/concrete.git
synced 2026-04-17 03:00:54 -04:00
d739e6672d
- this input index will be useful for MLIR/lower level conversions - it represents the input index of an Input node when considering the traced function signature - update code preparing function parameters to keep signature order
46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
"""hnumpy tracing utilities"""
|
|
from typing import Callable, Dict
|
|
|
|
import networkx as nx
|
|
|
|
from ..common.data_types import BaseValue
|
|
from ..common.tracing import (
|
|
BaseTracer,
|
|
create_graph_from_output_tracers,
|
|
make_input_tracers,
|
|
prepare_function_parameters,
|
|
)
|
|
|
|
|
|
class NPTracer(BaseTracer):
|
|
"""Tracer class for numpy operations"""
|
|
|
|
|
|
def trace_numpy_function(
|
|
function_to_trace: Callable, function_parameters: Dict[str, BaseValue]
|
|
) -> nx.MultiDiGraph:
|
|
"""Function used to trace a numpy function
|
|
|
|
Args:
|
|
function_to_trace (Callable): The function you want to trace
|
|
function_parameters (Dict[str, BaseValue]): A dictionary indicating what each input of the
|
|
function is e.g. an EncryptedValue holding a 7bits unsigned Integer
|
|
|
|
Returns:
|
|
nx.MultiDiGraph: The graph containing the ir nodes representing the computation done in the
|
|
input function
|
|
"""
|
|
function_parameters = prepare_function_parameters(function_to_trace, function_parameters)
|
|
|
|
input_tracers = make_input_tracers(NPTracer, function_parameters)
|
|
|
|
# We could easily create a graph of NPTracer, but we may end up with dead nodes starting from
|
|
# the inputs that's why we create the graph starting from the outputs
|
|
output_tracers = function_to_trace(**input_tracers)
|
|
if isinstance(output_tracers, NPTracer):
|
|
output_tracers = (output_tracers,)
|
|
|
|
graph = create_graph_from_output_tracers(output_tracers)
|
|
|
|
return graph
|