peft.LoRADenseGeneralAdapter

class gemma.peft.LoRADenseGeneralAdapter(*, rank: int, features: int | collections.abc.Sequence[int], axis: int | collections.abc.Sequence[int], batch_dims: collections.abc.Sequence[int], dtype: numpy.dtype = <class 'jax.numpy.float64'>, a_init: jax.nn.initializers.Initializer | collections.abc.Callable[[...], typing.Any] = <function variance_scaling.<locals>.init>, b_init: jax.nn.initializers.Initializer | collections.abc.Callable[[...], typing.Any] = <function zeros>, parent: flax.linen.module.Module | flax.core.scope.Scope | flax.linen.module._Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None)

Bases: flax.linen.module.Module

LoRA general dense module.

This module only does the x @ A @ B computation. Use LoRAGeneralDense to wrap a nn.Einsum layer.

rank

The rank of the LoRA decomposition.

Type:

int

features

The number of output features.

Type:

int | collections.abc.Sequence[int]

axis

int or tuple with axes to apply the transformation on.

Type:

int | collections.abc.Sequence[int]

batch_dims

tuple with batch axes.

Type:

collections.abc.Sequence[int]

dtype

The dtype to use for the LoRA weights.

Type:

numpy.dtype

a_init

The initializer for the A matrix.

Type:

jax.nn.initializers.Initializer | collections.abc.Callable[[…], Any]

b_init

The initializer for the B matrix.

Type:

jax.nn.initializers.Initializer | collections.abc.Callable[[…], Any]

rank: int

features: int | collections.abc.Sequence[int]

axis: int | collections.abc.Sequence[int]

batch_dims: collections.abc.Sequence[int]

dtype

alias of jax.numpy.float64

a_init(

shape: collections.abc.Sequence[int | Any],

dtype: Any | None = None,

out_sharding: jax.sharding.NamedSharding | jax.P | None = None,

) → jax.Array

b_init(

shape: collections.abc.Sequence[int | Any],

dtype: Any | None = None,

out_sharding: jax.sharding.NamedSharding | jax.P | None = None,

) → jax.Array

An initializer that returns a constant array full of zeros.

The key argument is ignored.

>>> import jax, jax.numpy as jnp
>>> jax.nn.initializers.zeros(jax.random.key(42), (2, 3), jnp.float32)
Array([[0., 0., 0.],
       [0., 0., 0.]], dtype=float32)

name: str | None = None

parent: flax.linen.module.Module | flax.core.scope.Scope | flax.linen.module._Sentinel | None = None

scope: flax.core.scope.Scope | None = None