ControlDependency

Insert dependencies between arrays in the graph.

Abstract Signature:

ControlDependency(inputs: Union[Tensor, Sequence[Tensor]], dependencies: Union[Tensor, Sequence[Tensor]])

JAX (Core)TensorFlowApple MLXFlax NNXPaxML / Praxis

JAX (Core)

API: —

Strategy: Plugin (jax_control_dependencies)

TensorFlow

API: tf.control_dependencies

Strategy: Direct Mapping

Official Docs ↗

Apple MLX

API: mlx.core.depends

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: —

Strategy: Plugin (jax_control_dependencies)

PaxML / Praxis

API: —

Strategy: Plugin (jax_control_dependencies)