// Welcome to JAX-js REPL!
// JAX-like numerical computing for the web
// Running on WebGPU & WebAssembly

import { numpy as np } from "@jax-js/jax";

// Create arrays
const x = np.array([1, 2, 3, 4, 5]);
const y = np.array([10, 20, 30, 40, 50]);

// Mathematical operations
const sum = x.ref.add(y.ref);
const product = x.ref.mul(y.ref);

// Print results
console.log("x:", await x.ref.jsAsync());
console.log("y:", await y.ref.jsAsync());
console.log("x + y:", await sum.ref.jsAsync());
console.log("x * y:", await product.ref.jsAsync());

// Cleanup
x.dispose();
y.dispose();
sum.dispose();
product.dispose();

return "Computation complete!";