我有一個資料類，如：
import dataclasses
import jax.numpy as jnp
@dataclasses.dataclass
class Metric:
score1: jnp.ndarray
score2: jnp.ndarray
score3: jnp.ndarray
在我的代碼中，我創建了它的多個實體，是否有一種簡單的方法可以合并每個屬性中的兩個屬性？例如，如果我有：
a = Metric(score1=jnp.array([10,10,10]), score2=jnp.array([20,20,20]), score3=jnp.array([30,30,30]))
b = Metric(score1=jnp.array([10,10,10]), score2=jnp.array([20,20,20]), score3=jnp.array([30,30,30]))
我想合并它們，例如擁有一個包含：
score1=jnp.array([10,10,10,10,10,10]), score2=jnp.array([20,20,20,20,20,20]) and score3=jnp.array([30,30,30,30,30,30])
uj5u.com熱心網友回復：
最簡單的事情可能只是定義一個方法：
import dataclasses
import jax.numpy as jnp
@dataclasses.dataclass
class Metric:
score1: jnp.ndarray
score2: jnp.ndarray
score3: jnp.ndarray
def concatenate(self, other):
return Metric(
jnp.concatenate((self.score1, other.score1)),
jnp.concatenate((self.score2, other.score2)),
jnp.concatenate((self.score3, other.score3)),
)
然后就做a.concatenate(b)。您也可以改為呼叫該方法__add__，這樣就可以使用a b. 這更整潔，但可能會與元素添加混淆。
uj5u.com熱心網友回復：
通過將類注冊Metric為pytree_node. google/flax，一個建立在 之上的神經網路庫jax，提供了flax.struct.dataclass這樣做的幫助。注冊后，您可以使用該jax.tree_util包來操作Metric實體：
from flax.struct import dataclass as flax_dataclass
from jax.tree_util import tree_multimap
import jax.numpy as jnp
@flax_dataclass
class Metric:
score1: jnp.ndarray
score2: jnp.ndarray
score3: jnp.ndarray
a = Metric(score1=jnp.array([10,10,10]), score2=jnp.array([20,20,20]), score3=jnp.array([30,30,30]))
b = Metric(score1=jnp.array([10,10,10]), score2=jnp.array([20,20,20]), score3=jnp.array([30,30,30]))
tree_multimap(lambda x, y: jnp.concatenate([x, y]), a, b)
給出：
Metric(score1=DeviceArray([10, 10, 10, 10, 10, 10], dtype=int32), score2=DeviceArray([20, 20, 20, 20, 20, 20], dtype=int32), score3=DeviceArray([30, 30, 30, 30, 30, 30], dtype=int32))
轉載請註明出處，本文鏈接：https://www.uj5u.com/qiye/424172.html
下一篇：將兩個列分組后求平方和