mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix Initializer protocol

Browse Source
This commit is contained in:
Klaus Greff 2025-02-26 14:25:15 +01:00 committed by GitHub
parent 99a12ef9ea
commit 5acfc88a00
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
jax/_src/nn/initializers.py
View File

@ -46,8 +46,8 @@ RealNumeric = Any # Scalar jnp array or float
@export
@typing.runtime_checkable
class Initializer(Protocol):
@staticmethod
def __call__(key: Array,
def __call__(self,
key: Array,
shape: core.Shape,
dtype: DTypeLikeInexact = jnp.float_) -> Array:
raise NotImplementedError