EvalShape

Computes the shape/dtype of a function without performing actual computation.

Abstract Signature:

EvalShape(f: Callable, args)

PyTorch

API: torch.func.eval_shape
Strategy: Direct Mapping
Official Docs ↗

JAX (Core)

API: jax.eval_shape
Strategy: Direct Mapping
Official Docs ↗

TensorFlow

API: —
Strategy: Custom / Partial

Flax NNX

API: flax.nnx.eval_shape
Strategy: Direct Mapping
Official Docs ↗