nbxtensor 0.0.1__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.
Files changed (291) hide show
  1. nbxtensor/THIRD_PARTY_NOTICES.md +287 -0
  2. nbxtensor/__init__.py +17 -0
  3. nbxtensor/backend/__init__.py +1 -0
  4. nbxtensor/backend/capabilities.py +207 -0
  5. nbxtensor/backend/cpu/__init__.py +1 -0
  6. nbxtensor/backend/cpu/blas.py +102 -0
  7. nbxtensor/backend/cpu/pool.py +204 -0
  8. nbxtensor/backend/cpu_backend.py +162 -0
  9. nbxtensor/backend/device_facts.py +105 -0
  10. nbxtensor/backend/device_files.py +129 -0
  11. nbxtensor/backend/device_names.py +18 -0
  12. nbxtensor/backend/devices/amd/cdna.yml +228 -0
  13. nbxtensor/backend/devices/amd/cdna2.yml +228 -0
  14. nbxtensor/backend/devices/amd/cdna3.yml +227 -0
  15. nbxtensor/backend/devices/apple/apple_m1.yml +266 -0
  16. nbxtensor/backend/devices/apple/apple_m1_max.yml +266 -0
  17. nbxtensor/backend/devices/apple/apple_m1_pro.yml +266 -0
  18. nbxtensor/backend/devices/apple/apple_m1_ultra.yml +266 -0
  19. nbxtensor/backend/devices/apple/apple_m2.yml +266 -0
  20. nbxtensor/backend/devices/apple/apple_m2_max.yml +266 -0
  21. nbxtensor/backend/devices/apple/apple_m2_pro.yml +266 -0
  22. nbxtensor/backend/devices/apple/apple_m2_ultra.yml +266 -0
  23. nbxtensor/backend/devices/apple/apple_m3.yml +266 -0
  24. nbxtensor/backend/devices/apple/apple_m3_max.yml +266 -0
  25. nbxtensor/backend/devices/apple/apple_m3_pro.yml +266 -0
  26. nbxtensor/backend/devices/apple/apple_m3_ultra.yml +266 -0
  27. nbxtensor/backend/devices/apple/apple_m4.yml +266 -0
  28. nbxtensor/backend/devices/apple/apple_m4_max.yml +266 -0
  29. nbxtensor/backend/devices/apple/apple_m4_pro.yml +308 -0
  30. nbxtensor/backend/devices/apple/apple_m5.yml +269 -0
  31. nbxtensor/backend/devices/apple/apple_m5_max.yml +269 -0
  32. nbxtensor/backend/devices/apple/apple_m5_pro.yml +269 -0
  33. nbxtensor/backend/devices/apple/apple_m5_ultra.yml +269 -0
  34. nbxtensor/backend/devices/apple/apple_m6.yml +269 -0
  35. nbxtensor/backend/devices/apple/apple_silicon.yml +263 -0
  36. nbxtensor/backend/devices/nvidia/ampere.yml +214 -0
  37. nbxtensor/backend/devices/nvidia/hopper.yml +213 -0
  38. nbxtensor/backend/devices/nvidia/volta.yml +411 -0
  39. nbxtensor/backend/faults.py +151 -0
  40. nbxtensor/backend/host_memory.py +185 -0
  41. nbxtensor/backend/launcher.py +2350 -0
  42. nbxtensor/backend/launcher_contract.py +429 -0
  43. nbxtensor/backend/metal_backend.py +661 -0
  44. nbxtensor/backend/metal_device.py +1274 -0
  45. nbxtensor/backend/process_memory.py +73 -0
  46. nbxtensor/backend/replay_graph.py +322 -0
  47. nbxtensor/backend/runtime.py +409 -0
  48. nbxtensor/backend/triton_ext_driver.py +735 -0
  49. nbxtensor/dtype/__init__.py +1 -0
  50. nbxtensor/dtype/convert.py +107 -0
  51. nbxtensor/dtype/dtypes.yml +70 -0
  52. nbxtensor/dtype/itemsize.py +166 -0
  53. nbxtensor/dtype/nbx_dtype.py +72 -0
  54. nbxtensor/dtype/operands.py +158 -0
  55. nbxtensor/dtype/outputs.py +64 -0
  56. nbxtensor/dtype/policy.py +155 -0
  57. nbxtensor/dtype/policy.yml +51 -0
  58. nbxtensor/dtype/promotion.py +112 -0
  59. nbxtensor/host.py +55 -0
  60. nbxtensor/io/__init__.py +1 -0
  61. nbxtensor/io/read_rate.py +197 -0
  62. nbxtensor/io/safetensors.py +314 -0
  63. nbxtensor/memory/__init__.py +9 -0
  64. nbxtensor/memory/allocator.py +1988 -0
  65. nbxtensor/memory/arena.py +75 -0
  66. nbxtensor/memory/device_memory.py +150 -0
  67. nbxtensor/memory/device_transfer.py +154 -0
  68. nbxtensor/memory/layout.py +43 -0
  69. nbxtensor/memory/memory_pool.py +191 -0
  70. nbxtensor/memory/metal.py +96 -0
  71. nbxtensor/memory/oom_advice.py +91 -0
  72. nbxtensor/memory/pointer_tables.py +323 -0
  73. nbxtensor/memory/pool.py +507 -0
  74. nbxtensor/memory/shadow.py +43 -0
  75. nbxtensor/memory/slab.py +273 -0
  76. nbxtensor/memory/trace_replay.py +122 -0
  77. nbxtensor/memory/traces/sana-video-2b-720p.16g.2026-11-02.tsv +13610 -0
  78. nbxtensor/memory/vmm.py +118 -0
  79. nbxtensor/ops/__init__.py +1 -0
  80. nbxtensor/ops/_autotune_policy.py +90 -0
  81. nbxtensor/ops/_common.py +32 -0
  82. nbxtensor/ops/_configs.py +1041 -0
  83. nbxtensor/ops/_cost.py +231 -0
  84. nbxtensor/ops/abs.py +17 -0
  85. nbxtensor/ops/adaptive_avg_pool2d.py +161 -0
  86. nbxtensor/ops/add.py +96 -0
  87. nbxtensor/ops/addcdiv.py +29 -0
  88. nbxtensor/ops/addcmul.py +29 -0
  89. nbxtensor/ops/addmv_op.py +64 -0
  90. nbxtensor/ops/addr.py +68 -0
  91. nbxtensor/ops/all_reduce.py +83 -0
  92. nbxtensor/ops/amax.py +40 -0
  93. nbxtensor/ops/any_reduce.py +83 -0
  94. nbxtensor/ops/arange_op.py +27 -0
  95. nbxtensor/ops/argmax.py +107 -0
  96. nbxtensor/ops/argmin.py +93 -0
  97. nbxtensor/ops/atan2.py +71 -0
  98. nbxtensor/ops/attention_wrappers.py +1648 -0
  99. nbxtensor/ops/autotune_bucket.py +137 -0
  100. nbxtensor/ops/autotune_cache.py +430 -0
  101. nbxtensor/ops/autotune_certified.py +1262 -0
  102. nbxtensor/ops/autotune_refusals.py +380 -0
  103. nbxtensor/ops/autotune_settings.py +58 -0
  104. nbxtensor/ops/avg_pool2d.py +73 -0
  105. nbxtensor/ops/baddbmm_op.py +132 -0
  106. nbxtensor/ops/batch_norm.py +163 -0
  107. nbxtensor/ops/binary_strided.py +196 -0
  108. nbxtensor/ops/bitwise_and.py +26 -0
  109. nbxtensor/ops/bitwise_left_shift.py +47 -0
  110. nbxtensor/ops/bitwise_not.py +24 -0
  111. nbxtensor/ops/bitwise_or.py +26 -0
  112. nbxtensor/ops/bitwise_right_shift.py +47 -0
  113. nbxtensor/ops/bucketize.py +37 -0
  114. nbxtensor/ops/cat_op.py +87 -0
  115. nbxtensor/ops/celu.py +20 -0
  116. nbxtensor/ops/clamp.py +24 -0
  117. nbxtensor/ops/clamp_min.py +22 -0
  118. nbxtensor/ops/conv2d.py +265 -0
  119. nbxtensor/ops/conv2d_m8n8k4.py +165 -0
  120. nbxtensor/ops/conv_depthwise2d.py +154 -0
  121. nbxtensor/ops/conv_transpose2d.py +148 -0
  122. nbxtensor/ops/conv_wrappers.py +1482 -0
  123. nbxtensor/ops/copy.py +18 -0
  124. nbxtensor/ops/copy_op.py +39 -0
  125. nbxtensor/ops/cos.py +17 -0
  126. nbxtensor/ops/count_nonzero.py +64 -0
  127. nbxtensor/ops/cross_entropy.py +67 -0
  128. nbxtensor/ops/cummax.py +109 -0
  129. nbxtensor/ops/cummin.py +105 -0
  130. nbxtensor/ops/cumsum.py +244 -0
  131. nbxtensor/ops/decode_attn_vec.py +162 -0
  132. nbxtensor/ops/depthwise_conv2d.py +132 -0
  133. nbxtensor/ops/dequant_gemv.py +190 -0
  134. nbxtensor/ops/dft.py +67 -0
  135. nbxtensor/ops/dispatch.py +1228 -0
  136. nbxtensor/ops/div.py +75 -0
  137. nbxtensor/ops/dot_op.py +67 -0
  138. nbxtensor/ops/dropout.py +23 -0
  139. nbxtensor/ops/dtype_convert.py +42 -0
  140. nbxtensor/ops/elementwise_launch.py +81 -0
  141. nbxtensor/ops/elementwise_wrappers.py +2315 -0
  142. nbxtensor/ops/elu.py +20 -0
  143. nbxtensor/ops/embedding.py +69 -0
  144. nbxtensor/ops/eq.py +32 -0
  145. nbxtensor/ops/erf.py +18 -0
  146. nbxtensor/ops/exp.py +17 -0
  147. nbxtensor/ops/exp2.py +18 -0
  148. nbxtensor/ops/fft_op.py +125 -0
  149. nbxtensor/ops/fill_op.py +20 -0
  150. nbxtensor/ops/flash_attention.py +307 -0
  151. nbxtensor/ops/flash_attention_m8n8k4.py +384 -0
  152. nbxtensor/ops/flash_decode.py +78 -0
  153. nbxtensor/ops/flip_op.py +60 -0
  154. nbxtensor/ops/floor.py +90 -0
  155. nbxtensor/ops/full_op.py +24 -0
  156. nbxtensor/ops/fused_moe.py +417 -0
  157. nbxtensor/ops/fused_moe_m8n8k4.py +99 -0
  158. nbxtensor/ops/fused_upsample_conv.py +413 -0
  159. nbxtensor/ops/gather_op.py +52 -0
  160. nbxtensor/ops/ge.py +32 -0
  161. nbxtensor/ops/gelu.py +31 -0
  162. nbxtensor/ops/gemv_vec.py +80 -0
  163. nbxtensor/ops/glu.py +43 -0
  164. nbxtensor/ops/grid_sampler.py +145 -0
  165. nbxtensor/ops/group_norm_tiled.py +67 -0
  166. nbxtensor/ops/groupnorm.py +123 -0
  167. nbxtensor/ops/gt.py +32 -0
  168. nbxtensor/ops/hardsigmoid.py +18 -0
  169. nbxtensor/ops/hardswish.py +19 -0
  170. nbxtensor/ops/im2col.py +57 -0
  171. nbxtensor/ops/index_add.py +95 -0
  172. nbxtensor/ops/index_put_op.py +92 -0
  173. nbxtensor/ops/index_select.py +121 -0
  174. nbxtensor/ops/index_wrappers.py +770 -0
  175. nbxtensor/ops/isfinite.py +20 -0
  176. nbxtensor/ops/isinf.py +20 -0
  177. nbxtensor/ops/isnan.py +18 -0
  178. nbxtensor/ops/launch_keys.py +1054 -0
  179. nbxtensor/ops/launch_observer.py +54 -0
  180. nbxtensor/ops/layernorm.py +71 -0
  181. nbxtensor/ops/le.py +35 -0
  182. nbxtensor/ops/leaky_relu.py +19 -0
  183. nbxtensor/ops/lerp.py +76 -0
  184. nbxtensor/ops/linalg_wrappers.py +1370 -0
  185. nbxtensor/ops/linspace_op.py +47 -0
  186. nbxtensor/ops/log.py +17 -0
  187. nbxtensor/ops/log_sigmoid.py +23 -0
  188. nbxtensor/ops/logical_and.py +26 -0
  189. nbxtensor/ops/logical_not.py +24 -0
  190. nbxtensor/ops/logical_or.py +26 -0
  191. nbxtensor/ops/logical_xor.py +37 -0
  192. nbxtensor/ops/loss_wrappers.py +282 -0
  193. nbxtensor/ops/lt.py +32 -0
  194. nbxtensor/ops/masked_fill.py +21 -0
  195. nbxtensor/ops/matmul.py +377 -0
  196. nbxtensor/ops/matmul_m8n8k4.py +235 -0
  197. nbxtensor/ops/max_pool2d.py +74 -0
  198. nbxtensor/ops/max_pool3d.py +93 -0
  199. nbxtensor/ops/max_reduce.py +91 -0
  200. nbxtensor/ops/maximum.py +18 -0
  201. nbxtensor/ops/mean.py +40 -0
  202. nbxtensor/ops/min_reduce.py +91 -0
  203. nbxtensor/ops/minimum.py +18 -0
  204. nbxtensor/ops/mish.py +22 -0
  205. nbxtensor/ops/mma_m8n8k4.py +152 -0
  206. nbxtensor/ops/moe_align.py +195 -0
  207. nbxtensor/ops/moe_decode_vec.py +374 -0
  208. nbxtensor/ops/moe_device_residency.py +141 -0
  209. nbxtensor/ops/moe_wrappers.py +345 -0
  210. nbxtensor/ops/mse_loss.py +79 -0
  211. nbxtensor/ops/mul.py +51 -0
  212. nbxtensor/ops/nan_to_num.py +23 -0
  213. nbxtensor/ops/ne.py +32 -0
  214. nbxtensor/ops/neg.py +17 -0
  215. nbxtensor/ops/nllloss.py +124 -0
  216. nbxtensor/ops/nonzero.py +33 -0
  217. nbxtensor/ops/nonzero_op.py +62 -0
  218. nbxtensor/ops/norm_wrappers.py +579 -0
  219. nbxtensor/ops/ones_op.py +20 -0
  220. nbxtensor/ops/oracles/__init__.py +6 -0
  221. nbxtensor/ops/oracles/conv2d_fp64.py +450 -0
  222. nbxtensor/ops/pad_op.py +159 -0
  223. nbxtensor/ops/paged_decode_attn.py +215 -0
  224. nbxtensor/ops/pixel_shuffle.py +144 -0
  225. nbxtensor/ops/pixel_unshuffle.py +72 -0
  226. nbxtensor/ops/pow.py +137 -0
  227. nbxtensor/ops/prod.py +81 -0
  228. nbxtensor/ops/rand_op.py +52 -0
  229. nbxtensor/ops/random_wrappers.py +366 -0
  230. nbxtensor/ops/reciprocal.py +17 -0
  231. nbxtensor/ops/reduction_wrappers.py +845 -0
  232. nbxtensor/ops/relu.py +18 -0
  233. nbxtensor/ops/remainder.py +46 -0
  234. nbxtensor/ops/repeat_interleave.py +48 -0
  235. nbxtensor/ops/repeat_op.py +132 -0
  236. nbxtensor/ops/residual_chain.py +374 -0
  237. nbxtensor/ops/rmsnorm.py +68 -0
  238. nbxtensor/ops/rope.py +142 -0
  239. nbxtensor/ops/rsqrt.py +17 -0
  240. nbxtensor/ops/scaled_softmax.py +96 -0
  241. nbxtensor/ops/scatter_op.py +139 -0
  242. nbxtensor/ops/screen_oracle.py +479 -0
  243. nbxtensor/ops/selu.py +22 -0
  244. nbxtensor/ops/shape_wrappers.py +1032 -0
  245. nbxtensor/ops/sigmoid.py +19 -0
  246. nbxtensor/ops/silu.py +27 -0
  247. nbxtensor/ops/sin.py +17 -0
  248. nbxtensor/ops/softmax.py +169 -0
  249. nbxtensor/ops/softplus.py +27 -0
  250. nbxtensor/ops/sort_op.py +304 -0
  251. nbxtensor/ops/spectral_wrappers.py +481 -0
  252. nbxtensor/ops/sqrt.py +17 -0
  253. nbxtensor/ops/stack_op.py +84 -0
  254. nbxtensor/ops/std.py +113 -0
  255. nbxtensor/ops/strided_copy.py +200 -0
  256. nbxtensor/ops/sub.py +35 -0
  257. nbxtensor/ops/sum.py +40 -0
  258. nbxtensor/ops/tan.py +19 -0
  259. nbxtensor/ops/tanh.py +21 -0
  260. nbxtensor/ops/threshold.py +19 -0
  261. nbxtensor/ops/tiling.yml +91 -0
  262. nbxtensor/ops/tiling_policy.py +43 -0
  263. nbxtensor/ops/tiling_sizes.py +370 -0
  264. nbxtensor/ops/topk.py +267 -0
  265. nbxtensor/ops/trace_op.py +64 -0
  266. nbxtensor/ops/tril.py +64 -0
  267. nbxtensor/ops/triu.py +62 -0
  268. nbxtensor/ops/unfold_op.py +37 -0
  269. nbxtensor/ops/unique_op.py +97 -0
  270. nbxtensor/ops/upsample_bilinear2d.py +106 -0
  271. nbxtensor/ops/upsample_nearest2d.py +52 -0
  272. nbxtensor/ops/upsample_nearest3d.py +61 -0
  273. nbxtensor/ops/var.py +136 -0
  274. nbxtensor/ops/vector_norm.py +216 -0
  275. nbxtensor/ops/weight_norm.py +104 -0
  276. nbxtensor/ops/where.py +39 -0
  277. nbxtensor/ops/wrappers.py +825 -0
  278. nbxtensor/ops/zeros_op.py +20 -0
  279. nbxtensor/quant/__init__.py +1 -0
  280. nbxtensor/quant/quantized_tensor.py +187 -0
  281. nbxtensor/random/__init__.py +1 -0
  282. nbxtensor/random/mt_randn.py +154 -0
  283. nbxtensor/random/rng_pin.py +59 -0
  284. nbxtensor/random/rng_stream.py +118 -0
  285. nbxtensor/random/seeded_draw.py +223 -0
  286. nbxtensor/shape/__init__.py +157 -0
  287. nbxtensor/tensor.py +1867 -0
  288. nbxtensor-0.0.1.dist-info/METADATA +379 -0
  289. nbxtensor-0.0.1.dist-info/RECORD +291 -0
  290. nbxtensor-0.0.1.dist-info/WHEEL +4 -0
  291. nbxtensor-0.0.1.dist-info/licenses/LICENSE +189 -0

There are too many changes on this page to be displayed.

The amount of changes on this page would crash your brower.

You can still verify the content by downloading the package file manually.