Add support for `dtype` / `DTypePolicy` to `JaxLayer` and `FlaxLayer`.#20732

fchollet merged 1 commit intokeras-team:masterkeras-team/keras:masterfrom

hertschuh:jax_layer_dtypehertschuh/keras:jax_layer_dtypeCopy head branch name to clipboard