Add `jaxtyping` for torch/jax code.

sustainability-lab / ASTRA

"AI for Sustainability" Toolkit for Research and Analysis

2 stars 7 forks source link

Open patel-zeel opened 10 months ago

patel-zeel commented 10 months ago

This blog explains the advanced usage of jaxtyping.

Actively raise an error if the shapes of inputs do not match with defined shapes.
Check if output_shape is input_shape - 2. It evaluates the math.