FlaxLayer

Keras Layer wrapping Flax.

Abstract Signature:

FlaxLayer(module)

Keras

API: keras.layers.FlaxLayer
Strategy: Direct Mapping
Official Docs ↗

TensorFlow

API: keras.layers.FlaxLayer
Strategy: Direct Mapping
Official Docs ↗