rejuvyesh/stresstest_jax_dlpack.jl

## stresstest_jax_dlpack.jl
using PyCall
using CUDA
using DLPack
using Test
#using Zygote
#using ChainRulesCore

@show DLPack.PYCALL_NOOP_DELETER

jax = pyimport("jax")
dlpack = pyimport("jax.dlpack")
numpy = pyimport("numpy")

pyto_dlpack(x) = @pycall dlpack.to_dlpack(x)::PyObject
pyfrom_dlpack(x) = @pycall dlpack.from_dlpack(x)::PyObject

@testset "dlpack" begin
    key = jax.random.PRNGKey(0)
    for dims in ((10,), (1, 10), (2, 3, 5), (2, 3, 4, 5))
        xto = jax.random.normal(key, dims)
        xjl = DLArray(xto, pyto_dlpack)
        @test isapprox(numpy.array(xto), Array(xjl))
    end
end

cujl = CuArray(randn(Float32, 2, 3, 5))
cujax = DLPack.share(cujl, pyfrom_dlpack)
@test (jax.numpy.sum(cujax).item()) ≈ sum(cujl)
	using PyCall
	using CUDA
	using DLPack
	using Test
	#using Zygote
	#using ChainRulesCore

	@show DLPack.PYCALL_NOOP_DELETER

	jax = pyimport("jax")
	dlpack = pyimport("jax.dlpack")
	numpy = pyimport("numpy")

	pyto_dlpack(x) = @pycall dlpack.to_dlpack(x)::PyObject
	pyfrom_dlpack(x) = @pycall dlpack.from_dlpack(x)::PyObject

	@testset "dlpack" begin
	key = jax.random.PRNGKey(0)
	for dims in ((10,), (1, 10), (2, 3, 5), (2, 3, 4, 5))
	xto = jax.random.normal(key, dims)
	xjl = DLArray(xto, pyto_dlpack)
	@test isapprox(numpy.array(xto), Array(xjl))
	end
	end

	cujl = CuArray(randn(Float32, 2, 3, 5))
	cujax = DLPack.share(cujl, pyfrom_dlpack)
	@test (jax.numpy.sum(cujax).item()) ≈ sum(cujl)