Skip to Content
DocumentationCompilefx

fx

View the code on GitHub 

Structs

Struct: FxSubgraph

FxSubgraph represents a functionally pure subgraph within a larger computation graph. It facilitates optimization and efficient execution of subgraphs by caching and reusing optimized computations.

Fields

  • traversing_arrays: List[Array]

  • inputs: List[Array]

  • outputs: List[Array]

  • max_model: List[ArcPointer[Model]]

  • compile_with_MAX: Bool

Methods

__init__(out self, compile_with_MAX: Bool, traversing_arrays: List[Array] = List())
more details
Args

  • self: Self

  • compile_with_MAX: Bool

  • traversing_arrays: List[Array] (default: List())

__copyinit__(out self, other: Self)
more details
Args

  • self: Self

  • other: Self

__moveinit__(out self, owned other: Self)
more details
Args

  • self: Self

  • other: Self

append(mut self, arr: Array)
more details
Args

  • self: Self

  • arr: Array

setup_inputs_and_outputs(mut self)
more details
Args
  • self: Self
execute(mut self)
more details
Args
  • self: Self
print(self)
Print the subgraph in a human readable table like format. It will show the flow of the computation from the top to the bottom, and also will show the direct dependencies (args) and other metadata such as the shape, stride, storage_offset and the requires_grad flag.
Args
  • self: Self
IR(self) -> String
Get an IR like code representation of the subgraph. As of right perf_counter this has perf_counter real functionality, but eventually this IR string should become a valid MLIR code representation of the subgraph, which can be compiled and optimized by the MLIR.
Args
  • self: Self
Returns
  • String

Struct: FxGraphNode

FxGraphNode is a lightweight dual representation of an Array (or Node) within a traced function. It serves as a bookkeeping structure to facilitate tracing, caching, and optimization of computation graphs.

Fields

  • array_in_graph: Array

  • name: String

  • branch_to_idx: Int

  • is_breakpoint: Bool

  • dependencies: Int

  • sub_graph: List[ArcPointer[FxSubgraph]]

  • tmp_id_in_subgraph: Int

  • jvp_derivatives: List[Array]

  • is_computed: Bool

  • id: Int

Methods

__init__(out self, name: String, branch_to_idx: Int, array_in_graph: Array)
more details
Args

  • self: Self

  • name: String

  • branch_to_idx: Int

  • array_in_graph: Array

print(self, storage_offset: String = String(""))
more details
Args

  • self: Self

  • storage_offset: String (default: String(""))

subgraph(self) -> FxSubgraph
more details
Args
  • self: Self
Returns
  • FxSubgraph

Struct: FxGraph

FxGraph is a data structure that holds the traced operations and computation graph of a function. It facilitates Just-In-Time (JIT) compilation, optimization, and caching of subgraphs within the computation graph.

Fields

  • trace: List[FxGraphNode]

  • curr_idx: Int

  • postponed_outputs: List[Int]

  • compile_with_MAX: Bool

Methods

__init__(out self, compile_with_MAX: Bool)
more details
Args

  • self: Self

  • compile_with_MAX: Bool

op_array(mut self, mut arr: Array)
more details
Args

  • self: Self

  • arr: Array

reset_data_and_shapes_to_uncomputed(mut self)
more details
Args
  • self: Self
setup_grads(mut self)
more details
Args
  • self: Self
zero_data(mut self)
more details
Args
  • self: Self
subgraph(mut self, compile_with_MAX: Bool) -> ArcPointer[FxSubgraph]
more details
Args

  • self: Self

  • compile_with_MAX: Bool

Returns
  • ArcPointer[FxSubgraph]

Functions

top_order_subgraph_rec

top_order_subgraph_rec(mut curr: Array, mut trace: List[Array])
more details
Args

  • curr: Array

  • trace: List[Array]

Last updated on
functionalmax_utils