tony/tensorneat-mend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix bugs

Browse Source
This commit is contained in:
wls2002
2024-06-20 16:32:52 +08:00
parent 9f72813c35
commit 075460f896
17 changed files with 224 additions and 140 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tensorneat/utils/aggregation/agg_jnp.py
View File

@@ -9,19 +9,19 @@ class Agg:
@staticmethod
def sum(z):
return jnp.sum(z, axis=0, where=~jnp.isnan(z))
return jnp.sum(z, axis=0, where=~jnp.isnan(z), initial=0)
@staticmethod
def product(z):
return jnp.prod(z, axis=0, where=~jnp.isnan(z))
return jnp.prod(z, axis=0, where=~jnp.isnan(z), initial=1)
@staticmethod
def max(z):
return jnp.max(z, axis=0, where=~jnp.isnan(z))
return jnp.max(z, axis=0, where=~jnp.isnan(z), initial=-jnp.inf)
@staticmethod
def min(z):
return jnp.min(z, axis=0, where=~jnp.isnan(z))
return jnp.min(z, axis=0, where=~jnp.isnan(z), initial=jnp.inf)
@staticmethod
def maxabs(z):