性能优化 #

import jax
import jax.numpy as jnp

@jax.jit
def compute(x):
    return jnp.dot(x, x.T)

x_float32 = jnp.ones((100, 100), dtype=jnp.float32)
x_float64 = jnp.ones((100, 100), dtype=jnp.float64)

compute(x_float32)
compute(x_float64)
compute(x_float32)

import jax
import jax.numpy as jnp

@jax.jit(static_argnums=(1,))
def process(x, mode):
    if mode == 'relu':
        return jnp.maximum(x, 0)
    else:
        return x

x = jnp.array([-1, 2, -3, 4])
print(process(x, 'relu'))
print(process(x, 'linear'))

import os

os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'
os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '0.8'

import jax

import jax
import jax.numpy as jnp

def large_forward(params, x):
    for w in params:
        x = jnp.dot(x, w)
        x = jax.nn.relu(x)
    return x

@jax.remat
def memory_efficient_forward(params, x):
    return large_forward(params, x)

params = [jax.random.normal(jax.random.PRNGKey(i), (1024, 1024)) * 0.01 
          for i in range(50)]
x = jax.random.normal(jax.random.PRNGKey(100), (32, 1024))

output = memory_efficient_forward(params, x)
print(f"输出形状: {output.shape}")

import jax
import jax.numpy as jnp

def single_process(x):
    return jnp.sum(x ** 2)

batch_process = jax.jit(jax.vmap(single_process))

batch_x = jax.random.normal(jax.random.PRNGKey(0), (10000, 100))
result = batch_process(batch_x)
print(f"结果形状: {result.shape}")

import jax
import jax.numpy as jnp

@jax.pmap
def parallel_compute(x):
    return jnp.sum(x ** 2)

num_devices = jax.device_count()
batch_x = jax.random.normal(jax.random.PRNGKey(0), (num_devices, 1000, 100))
result = parallel_compute(batch_x)
print(f"结果形状: {result.shape}")

import jax
import jax.numpy as jnp

with jax.profiler.trace('/tmp/jax_trace'):
    @jax.jit
    def compute(x):
        return jnp.dot(x, x.T)
    
    x = jax.random.normal(jax.random.PRNGKey(0), (1000, 1000))
    result = compute(x)
    result.block_until_ready()

import jax
import jax.numpy as jnp
import time

@jax.jit
def compute(x):
    return jnp.dot(x, x.T)

x = jax.random.normal(jax.random.PRNGKey(0), (1000, 1000))

compute(x).block_until_ready()

start = time.time()
for _ in range(100):
    result = compute(x)
    result.block_until_ready()
end = time.time()

print(f"平均时间: {(end - start) / 100:.4f}s")