jax
42dd8f75 - Speed up DynamicJaxprTracer.__init__.

Commit
163 days ago
Speed up DynamicJaxprTracer.__init__. * Don't create MutableQDDs if we don't need them. Make mutable_qdd optional to save an allocation most of the time. * Use Tracer.__init__() rather than super().__init__(). This is slightly faster in a profile. * Add __slots__ on JaxprStackFrame to avoid some dictionary lookups. PiperOrigin-RevId: 917955593
Author
Parents
Loading