What is the Flax equivalent of Tensorflow functions?

[marcvanzee]

[marcvanzee]

I'm listing some Tensorflow functions below for which users asked what the equivalent is in Flax.

TF Function	Flax/jax Function
tf.identity	lambda x: x
tf.name_scope()	"Create submodule" [1]
tf.print	host_callback.id_print
tf.keras.initializers.Constant	lambda rng: 1.0

[1] Putting anything inside a new Module would have that effect. You can also always give submodules a name explicitly with the name=... kwarg to the submodule constuctor.

[avital]

A quick comment on [1] above -- the patterns suggested works for submodules defined "in-line" inside a method annotated with @nn.compact. If, instead, if you define submodules within setup, then the name of the attribute would be the new name scope (this follows the same pattern as PyTorch), e.g.:

class Foo(nn.Module):
  def setup(self):
    self.bar = Baz()

In this case the Baz submodule would be defined similarly to having been in a tf name scope named "bar".