rwkv-ops 0.2.1__py3-none-any.whl → 0.2.2__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.

This version of rwkv-ops might be problematic. Click here for more details.

Files changed (8) hide show
  1. rwkv_ops/__init__.py +1 -1
  2. rwkv_ops/rwkv7_kernel/__init__.py +0 -6
  3. rwkv_ops/rwkv7_kernel/jax_op.py +4 -3
  4. {rwkv_ops-0.2.1.dist-info → rwkv_ops-0.2.2.dist-info}/METADATA +1 -1
  5. {rwkv_ops-0.2.1.dist-info → rwkv_ops-0.2.2.dist-info}/RECORD +8 -8
  6. {rwkv_ops-0.2.1.dist-info → rwkv_ops-0.2.2.dist-info}/LICENSE.txt +0 -0
  7. {rwkv_ops-0.2.1.dist-info → rwkv_ops-0.2.2.dist-info}/WHEEL +0 -0
  8. {rwkv_ops-0.2.1.dist-info → rwkv_ops-0.2.2.dist-info}/top_level.txt +0 -0
rwkv_ops/__init__.py CHANGED
@@ -1,4 +1,4 @@
1
- __version__ = "0.2.1"
1
+ __version__ = "0.2.2"
2
2
  import os
3
3
 
4
4
  KERNEL_TYPE = os.environ.get("KERNEL_TYPE", "triton")
rwkv_ops/rwkv7_kernel/__init__.py CHANGED
@@ -132,14 +132,8 @@ def get_generalized_delta_rule(HEAD_SIZE=64, KERNEL_TYPE="native"):
132
132
  USE_KERNEL = False
133
133
  elif keras.config.backend() == "jax":
134
134
  from jax.lib import xla_bridge
135
- import jax
136
135
  import os
137
- import logging
138
136
 
139
- logging.basicConfig(level=logging.ERROR)
140
- os.environ["TRITON_LOG_LEVEL"] = "ERROR" # 只显示错误级别的日志
141
- os.environ["TRITON_DISABLE_AUTOTUNE"] = "1" # 禁用自动调优日志
142
- os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2" # 禁用自动调优日志
143
137
  if (
144
138
  xla_bridge.get_backend().platform == "gpu"
145
139
  and KERNEL_TYPE.lower() == "triton"
rwkv_ops/rwkv7_kernel/jax_op.py CHANGED
@@ -14,10 +14,9 @@ from .jax_kernel.chunk_o_fwd import chunk_dplr_fwd_o
14
14
  from .jax_kernel.wy_fast_bwd import chunk_dplr_bwd_wy
15
15
  from .jax_kernel.wy_fast_fwd import prepare_wy_repr_fwd
16
16
  from .jax_kernel.cumsum import chunk_rwkv6_fwd_cumsum
17
-
17
+ from jax.ad_checkpoint import checkpoint_policies
18
18
  CHUNKSIZE = 16
19
19
 
20
-
21
20
  def chunk_dplr_fwd(
22
21
  q: jax.Array,
23
22
  k: jax.Array,
@@ -156,7 +155,7 @@ def chunk_dplr_fwd_jax(
156
155
  output_final_state=True,
157
156
  )
158
157
  cache = (r, k, v, a, b, gk, initial_state)
159
- return [o, state], cache
158
+ return (o, state), cache
160
159
 
161
160
 
162
161
  def chunk_dplr_bwd(
@@ -378,3 +377,5 @@ def generalized_delta_rule(
378
377
  if output_final_state:
379
378
  return jnp.asarray(o, DTYPE), final_state
380
379
  return jnp.asarray(o, DTYPE)
380
+
381
+
{rwkv_ops-0.2.1.dist-info → rwkv_ops-0.2.2.dist-info}/METADATA RENAMED
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: rwkv-ops
3
- Version: 0.2.1
3
+ Version: 0.2.2
4
4
  Home-page: https://github.com/pass-lin/rwkv_ops
5
5
  License: Apache 2.0
6
6
  Keywords: rwkv implement for multi backend
{rwkv_ops-0.2.1.dist-info → rwkv_ops-0.2.2.dist-info}/RECORD RENAMED
@@ -1,12 +1,12 @@
1
- rwkv_ops/__init__.py,sha256=FhqDH54a_vrLJQLTWMEM-QnETJG2z6l4YjMiLtcZt2E,843
1
+ rwkv_ops/__init__.py,sha256=R44XEpC4o1C-bUJy43YokONCD5Jxp-kBp3eyNZNqCIk,843
2
2
  rwkv_ops/rwkv6_kernel/__init__.py,sha256=_j6G_3fY8xPxrlZbgDT2ndX4IPiNJ4qjqIcdmNI_r9Q,4100
3
3
  rwkv_ops/rwkv6_kernel/jax_rwkv_kernel.py,sha256=WOzqfQQSHHMoWqm2kRz_BhtMzGYc5USJ26qaEwuARo4,30117
4
4
  rwkv_ops/rwkv6_kernel/ops_rwkv_kernel.py,sha256=otjfw5n6nf2YVpBIWIZjaCsxMyLXXwg-ma1ueXX-EdY,3274
5
5
  rwkv_ops/rwkv6_kernel/torch_rwkv_kernel.py,sha256=Q1uPMgaS21OEfQ8-sBDjaCUASMtkSOdN3OosEUsBp9U,12918
6
- rwkv_ops/rwkv7_kernel/__init__.py,sha256=GpwZ5dk7d5H6u-VSUAQ29KQVnCZdEwtKZ13_3kVJets,5888
6
+ rwkv_ops/rwkv7_kernel/__init__.py,sha256=20UIBG-nydjB72iboCJ7TDIwgXVF1aYu93sqApFBqc8,5557
7
7
  rwkv_ops/rwkv7_kernel/get_jax_devices_info.py,sha256=cMIaNED7d1PvYNSyq8wNI3G7wNvcgdUj9HWRBLuSVM8,6004
8
8
  rwkv_ops/rwkv7_kernel/get_torch_devices_info.py,sha256=ZL_rAM6lHB4nTOOU28Xm08qptfuIoijOMi_xwJG3KCo,7380
9
- rwkv_ops/rwkv7_kernel/jax_op.py,sha256=tyMxvk_EblDaGsePpxw3AhELvolp7LeE5NopUhKw1R0,9107
9
+ rwkv_ops/rwkv7_kernel/jax_op.py,sha256=l7HMVUqf3M9YmE1OADoin93HpVsCqm4rnrKoH9s6Dzg,9158
10
10
  rwkv_ops/rwkv7_kernel/native_keras_op.py,sha256=QPrXLbqw0chipQg_0jepRp2U19BYpBBFdKZWyaDNNoc,2488
11
11
  rwkv_ops/rwkv7_kernel/torch_op.py,sha256=jw_AvqshTAG4t9-MRqxFQNi_bTzxNbx3lwnMifPk8-8,14070
12
12
  rwkv_ops/rwkv7_kernel/jax_kernel/__init__.py,sha256=uHsf_1qrtRK62IvhLuzefHGPWpHXmw1p0tqmwlHcptk,346
@@ -40,8 +40,8 @@ rwkv_ops/rwkv7_kernel/triton_kernel/cumsum.py,sha256=pRp_z587PrnpgRVpi031IndyjVI
40
40
  rwkv_ops/rwkv7_kernel/triton_kernel/utils.py,sha256=TNGlkwGq4t-TOcdVBk_N_vHPLzMFTu_F0V-O1RprIO4,553
41
41
  rwkv_ops/rwkv7_kernel/triton_kernel/wy_fast_bwd.py,sha256=szaG11q_WmpyhXi6aVWwzizvflCh5wND8wGA_V8afzA,5479
42
42
  rwkv_ops/rwkv7_kernel/triton_kernel/wy_fast_fwd.py,sha256=jbb19DUTHENU2RIOv_T4m_W1eXMqdRqG0XevIkBOhI4,9438
43
- rwkv_ops-0.2.1.dist-info/LICENSE.txt,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
44
- rwkv_ops-0.2.1.dist-info/METADATA,sha256=_I9thMqiJMM1kOlwUH23DVkRTnBHVLyjHQ5-74OvpS4,8411
45
- rwkv_ops-0.2.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
46
- rwkv_ops-0.2.1.dist-info/top_level.txt,sha256=cVqoKE-WR_e2gHL87-6O4K1kG6-yTJGB2huyr6FmD2I,9
47
- rwkv_ops-0.2.1.dist-info/RECORD,,
43
+ rwkv_ops-0.2.2.dist-info/LICENSE.txt,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
44
+ rwkv_ops-0.2.2.dist-info/METADATA,sha256=bBugt-UogSOxR0YxdSxnoAEPB7wb7oHhTrzlCEY5h_8,8411
45
+ rwkv_ops-0.2.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
46
+ rwkv_ops-0.2.2.dist-info/top_level.txt,sha256=cVqoKE-WR_e2gHL87-6O4K1kG6-yTJGB2huyr6FmD2I,9
47
+ rwkv_ops-0.2.2.dist-info/RECORD,,