Class MyModule

gemma/peft/_interceptors_test.py:44–61

class MyModule(nn.Module):

  @nn.compact
  def __call__(self, x):
    # Call a same model twice
    model = nn.Dense(1)
    y0 = nn.Dropout(0.5, deterministic=False)(x)
    y1 = model(x)
    y2 = model(x)

    # Call another model
    y3 = nn.Dense(2)(x)
    return {
        'y0': y0,
        'y1': y1,
        'y2': y2,
        'y3': y3,
    }


def test_module():