rwkv-ops 0.2__tar.gz → 0.2.2__tar.gz

{rwkv_ops-0.2 → rwkv_ops-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: rwkv_ops
-Version: 0.2
+Version: 0.2.2
 Home-page: https://github.com/pass-lin/rwkv_ops
 License: Apache 2.0
 Keywords: rwkv implement for multi backend

{rwkv_ops-0.2 → rwkv_ops-0.2.2}/rwkv_ops/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.2"
+__version__ = "0.2.2"
 import os
 
 KERNEL_TYPE = os.environ.get("KERNEL_TYPE", "triton")

{rwkv_ops-0.2 → rwkv_ops-0.2.2}/rwkv_ops/rwkv7_kernel/__init__.py

@@ -132,14 +132,8 @@ def get_generalized_delta_rule(HEAD_SIZE=64, KERNEL_TYPE="native"):
             USE_KERNEL = False
     elif keras.config.backend() == "jax":
         from jax.lib import xla_bridge
-        import jax
         import os
-        import logging
 
-        logging.basicConfig(level=logging.ERROR)
-        os.environ["TRITON_LOG_LEVEL"] = "ERROR"  # 只显示错误级别的日志
-        os.environ["TRITON_DISABLE_AUTOTUNE"] = "1"  # 禁用自动调优日志
-        os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"  # 禁用自动调优日志
         if (
             xla_bridge.get_backend().platform == "gpu"
             and KERNEL_TYPE.lower() == "triton"

{rwkv_ops-0.2 → rwkv_ops-0.2.2}/rwkv_ops/rwkv7_kernel/jax_op.py

@@ -14,10 +14,9 @@ from .jax_kernel.chunk_o_fwd import chunk_dplr_fwd_o
 from .jax_kernel.wy_fast_bwd import chunk_dplr_bwd_wy
 from .jax_kernel.wy_fast_fwd import prepare_wy_repr_fwd
 from .jax_kernel.cumsum import chunk_rwkv6_fwd_cumsum
-
+from jax.ad_checkpoint import checkpoint_policies 
 CHUNKSIZE = 16
 
-
 def chunk_dplr_fwd(
     q: jax.Array,
     k: jax.Array,
@@ -156,7 +155,7 @@ def chunk_dplr_fwd_jax(
         output_final_state=True,
     )
     cache = (r, k, v, a, b, gk, initial_state)
-    return [o, state], cache
+    return (o, state), cache
 
 
 def chunk_dplr_bwd(
@@ -378,3 +377,5 @@ def generalized_delta_rule(
     if output_final_state:
         return jnp.asarray(o, DTYPE), final_state
     return jnp.asarray(o, DTYPE)
+
+

{rwkv_ops-0.2 → rwkv_ops-0.2.2}/rwkv_ops/rwkv7_kernel/torch_op.py

@@ -473,16 +473,30 @@ def generalized_delta_rule(
     a = transpose_head(a, head_first)
     b = transpose_head(b, head_first)
     w = transpose_head(w, head_first)
-    out, state = chunk_rwkv7(
-        r=r,
-        k=k,
-        v=v,
-        a=a,
-        b=b,
-        w=w,
-        initial_state=initial_state,
-        output_final_state=output_final_state,
-    )
+    if w.device.type == "cuda":
+        out, state = chunk_rwkv7(
+            r=r,
+            k=k,
+            v=v,
+            a=a,
+            b=b,
+            w=w,
+            initial_state=initial_state,
+            output_final_state=output_final_state,
+        )
+    else:
+        from ops.native_keras_op import generalized_delta_rule
+
+        out, state = generalized_delta_rule(
+            r=r,
+            k=k,
+            v=v,
+            a=a,
+            b=b,
+            w=w,
+            initial_state=initial_state,
+            output_final_state=output_final_state,
+        )
     out = transpose_head(out, head_first)
     if output_final_state:
         return out, cast(state, dtype)

{rwkv_ops-0.2 → rwkv_ops-0.2.2}/rwkv_ops.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: rwkv-ops
-Version: 0.2
+Version: 0.2.2
 Home-page: https://github.com/pass-lin/rwkv_ops
 License: Apache 2.0
 Keywords: rwkv implement for multi backend

{rwkv_ops-0.2 → rwkv_ops-0.2.2}/setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="rwkv_ops",
-    version="0.2",
+    version="0.2.2",
     packages=find_packages(),
     install_requires=["keras"],  # 添加依赖项
     license="Apache 2.0",  # 指定许可证类型