@jax-js/jax 0.1.5 → 0.1.6

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.
@@ -151,6 +151,19 @@ function normalizeAxis(axis, ndim) {
151
151
  return sorted(seen);
152
152
  }
153
153
  }
154
+ /** Check for an array of integers with no duplicates. */
155
+ function checkInts(indices) {
156
+ if (typeof indices === "number") {
157
+ if (!Number.isInteger(indices)) throw new TypeError(`Expected integer index, got ${indices}`);
158
+ } else {
159
+ const seen = /* @__PURE__ */ new Set();
160
+ for (const i of indices) {
161
+ if (!Number.isInteger(i)) throw new TypeError(`Expected integer indices, got ${i}`);
162
+ if (seen.has(i)) throw new Error(`Duplicate index ${i} passed to function`);
163
+ seen.add(i);
164
+ }
165
+ }
166
+ }
154
167
  function range(start, stop, step = 1) {
155
168
  if (stop === void 0) {
156
169
  stop = start;
@@ -2306,10 +2319,10 @@ function tuneWebgpu(kernel) {
2306
2319
  upcastedAxis.add(choices[0][2]);
2307
2320
  } else break;
2308
2321
  }
2309
- if (/chrome/i.test(navigator.userAgent) && dim.reduce < dim.unroll && (prod(dim.st.shape.slice(dim.unroll)) <= 4 || dim.unroll === dim.upcast && prod(dim.st.shape.slice(dim.upcast)) < 64)) {
2322
+ if (!/Mobi|Android/i.test(navigator.userAgent) && dim.reduce < dim.unroll && (prod(dim.st.shape.slice(dim.unroll)) <= 4 || dim.unroll === dim.upcast && prod(dim.st.shape.slice(dim.upcast)) < 64)) {
2310
2323
  const s = dim.st.shape[dim.unroll - 1];
2311
2324
  if (0 < s && s <= 32) dim.applyUnroll(dim.reduce, s);
2312
- else for (const splits of [4]) if (s % splits === 0) {
2325
+ else for (const splits of [8, 4]) if (s % splits === 0) {
2313
2326
  dim.applyUnroll(dim.unroll - 1, splits);
2314
2327
  break;
2315
2328
  }
@@ -4240,7 +4253,7 @@ async function createBackend(device) {
4240
4253
  if (!navigator.gpu) return null;
4241
4254
  const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
4242
4255
  if (!adapter) return null;
4243
- const { WebGPUBackend } = await Promise.resolve().then(() => require("./webgpu-Db2JrNBr.cjs"));
4256
+ const { WebGPUBackend } = await Promise.resolve().then(() => require("./webgpu-rraa6dfz.cjs"));
4244
4257
  const importantLimits = [
4245
4258
  "maxBufferSize",
4246
4259
  "maxComputeInvocationsPerWorkgroup",
@@ -4278,7 +4291,7 @@ async function createBackend(device) {
4278
4291
  });
4279
4292
  if (!gl) return null;
4280
4293
  if (!gl.getExtension("EXT_color_buffer_float")) return null;
4281
- const { WebGLBackend } = await Promise.resolve().then(() => require("./webgl-ClIYb8jP.cjs"));
4294
+ const { WebGLBackend } = await Promise.resolve().then(() => require("./webgl-CyfzNW8T.cjs"));
4282
4295
  return new WebGLBackend(gl);
4283
4296
  } else throw new Error(`Backend not found: ${device}`);
4284
4297
  }
@@ -4446,6 +4459,12 @@ Object.defineProperty(exports, 'checkAxis', {
4446
4459
  return checkAxis;
4447
4460
  }
4448
4461
  });
4462
+ Object.defineProperty(exports, 'checkInts', {
4463
+ enumerable: true,
4464
+ get: function () {
4465
+ return checkInts;
4466
+ }
4467
+ });
4449
4468
  Object.defineProperty(exports, 'deepEqual', {
4450
4469
  enumerable: true,
4451
4470
  get: function () {
@@ -150,6 +150,19 @@ function normalizeAxis(axis, ndim) {
150
150
  return sorted(seen);
151
151
  }
152
152
  }
153
+ /** Check for an array of integers with no duplicates. */
154
+ function checkInts(indices) {
155
+ if (typeof indices === "number") {
156
+ if (!Number.isInteger(indices)) throw new TypeError(`Expected integer index, got ${indices}`);
157
+ } else {
158
+ const seen = /* @__PURE__ */ new Set();
159
+ for (const i of indices) {
160
+ if (!Number.isInteger(i)) throw new TypeError(`Expected integer indices, got ${i}`);
161
+ if (seen.has(i)) throw new Error(`Duplicate index ${i} passed to function`);
162
+ seen.add(i);
163
+ }
164
+ }
165
+ }
153
166
  function range(start, stop, step = 1) {
154
167
  if (stop === void 0) {
155
168
  stop = start;
@@ -2305,10 +2318,10 @@ function tuneWebgpu(kernel) {
2305
2318
  upcastedAxis.add(choices[0][2]);
2306
2319
  } else break;
2307
2320
  }
2308
- if (/chrome/i.test(navigator.userAgent) && dim.reduce < dim.unroll && (prod(dim.st.shape.slice(dim.unroll)) <= 4 || dim.unroll === dim.upcast && prod(dim.st.shape.slice(dim.upcast)) < 64)) {
2321
+ if (!/Mobi|Android/i.test(navigator.userAgent) && dim.reduce < dim.unroll && (prod(dim.st.shape.slice(dim.unroll)) <= 4 || dim.unroll === dim.upcast && prod(dim.st.shape.slice(dim.upcast)) < 64)) {
2309
2322
  const s = dim.st.shape[dim.unroll - 1];
2310
2323
  if (0 < s && s <= 32) dim.applyUnroll(dim.reduce, s);
2311
- else for (const splits of [4]) if (s % splits === 0) {
2324
+ else for (const splits of [8, 4]) if (s % splits === 0) {
2312
2325
  dim.applyUnroll(dim.unroll - 1, splits);
2313
2326
  break;
2314
2327
  }
@@ -4239,7 +4252,7 @@ async function createBackend(device) {
4239
4252
  if (!navigator.gpu) return null;
4240
4253
  const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
4241
4254
  if (!adapter) return null;
4242
- const { WebGPUBackend } = await import("./webgpu-Dh7k9io0.js");
4255
+ const { WebGPUBackend } = await import("./webgpu-C-VfevQW.js");
4243
4256
  const importantLimits = [
4244
4257
  "maxBufferSize",
4245
4258
  "maxComputeInvocationsPerWorkgroup",
@@ -4277,7 +4290,7 @@ async function createBackend(device) {
4277
4290
  });
4278
4291
  if (!gl) return null;
4279
4292
  if (!gl.getExtension("EXT_color_buffer_float")) return null;
4280
- const { WebGLBackend } = await import("./webgl-RSuZKvgc.js");
4293
+ const { WebGLBackend } = await import("./webgl-CLLvzJlO.js");
4281
4294
  return new WebGLBackend(gl);
4282
4295
  } else throw new Error(`Backend not found: ${device}`);
4283
4296
  }
@@ -4313,4 +4326,4 @@ var UnsupportedRoutineError = class extends Error {
4313
4326
  };
4314
4327
 
4315
4328
  //#endregion
4316
- export { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, Executable, FpHash, Kernel, PPrint, Reduction, Routine, Routines, ShapeTracker, SlotError, UnsupportedOpError, UnsupportedRoutineError, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, findPow2, generalBroadcast, getBackend, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, mapSetUnion, normalizeAxis, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, strip1, toposort, tuneNullopt, tuneWebgpu, unravelAlu, unzip2, zip, zipn };
4329
+ export { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, Executable, FpHash, Kernel, PPrint, Reduction, Routine, Routines, ShapeTracker, SlotError, UnsupportedOpError, UnsupportedRoutineError, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, checkInts, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, findPow2, generalBroadcast, getBackend, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, mapSetUnion, normalizeAxis, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, strip1, toposort, tuneNullopt, tuneWebgpu, unravelAlu, unzip2, zip, zipn };