orbax
orbax copied to clipboard
Support TensorSpecWithDefault in input_signature
Support TensorSpecWithDefault in input_signature