GleasonK

# NOTE: The JAX APIs for creating a composite are WIP, we are looking into a
# jax.lax.composite API -- this file as-is uses JAX internal APIs which are subject
# to change and should not be relied on for anything in production, but may be OK for
# experimentation.

"""XLA Composite Test."""

from absl.testing import absltest
import jax
from jax import export as jax_export

GleasonK

Example Log

The following is a log of the methods called to load the PJRT plugin and execute y = jax.jit(lambda x: jnp.power(x, jnp.int32(2)))(1).

I client_cpp_pjrt.cc:55] StablehloReferencePjrtClient(0x23bac400)
I device.cc:53] StablehloReferenceDeviceDescription(0x23bac4f8)
I device.cc:104] StablehloReferenceDevice(0x23bac4e0)
I device.cc:123] client(0x23bac4e0)

GleasonK

// This is a dump from PyTorch/XLA of the following LLAMA2 model file:
// https://github.com/pytorch/xla/blob/master/test/stablehlo/llama_model2.py
//
module @IrToHlo.7509 attributes {mhlo.cross_program_prefetches = [], mhlo.input_output_alias = [], mhlo.is_dynamic = false, mhlo.use_auto_spmd_partitioning = false} {
  func.func @main(%arg0: tensor<32000x4096xf32>, %arg1: tensor<4096xf32>, %arg2: tensor<4096x11008xf32>, %arg3: tensor<11008x4096xf32>, %arg4: tensor<4096xf32>, %arg5: tensor<4096x4096xf32>, %arg6: tensor<4096x4096xf32>, %arg7: tensor<4096xf32>, %arg8: tensor<4096x11008xf32>, %arg9: tensor<11008x4096xf32>, %arg10: tensor<4096xf32>, %arg11: tensor<4096x4096xf32>, %arg12: tensor<4096x4096xf32>, %arg13: tensor<4096xf32>, %arg14: tensor<4096x11008xf32>, %arg15: tensor<11008x4096xf32>, %arg16: tensor<4096xf32>, %arg17: tensor<4096x4096xf32>, %arg18: tensor<4096x4096xf32>, %arg19: tensor<4096xf32>, %arg20: tensor<4096x11008xf32>, %arg21: tensor<11008x4096xf32>, %arg22: tensor<4096xf32>, %arg23: tensor<4

GleasonK

// import jax
// import jax.numpy as jnp
// from jax.experimental.export import export
// 
// def f(a,b):
//   return jnp.add(a,b)
// 
// lhs = jax.ShapeDtypeStruct(export.symbolic_shape("a,10"), np.float32)
// rhs = jax.ShapeDtypeStruct(export.symbolic_shape("1"), np.float32)
// export.export(f)(polyA, poly2A)

GleasonK

# Useful script for detaching the screen of Surface Book.
# Run if detach light blinks red indicating that it cannot be detached
# since applications are preventing the detach.
#
# Option 1: Copy-paste into powershell window.
#
# Option 2: Run as PowerShell script (requires execution policy change):
# https://superuser.com/questions/106360/how-to-enable-execution-of-powershell-scripts
#
Try {

GleasonK

<div id="player"></div>
<div style="float: left; width: 50%;">
    <div id="video-chat" hidden="true" style="margin-bottom: 10px;">
		<div id="vid-box"></div>
		<button onclick="end()">End Call</button>
    </div>
    
    <form name="loginForm" id="login" action="#" onsubmit="return login(this);">
        	<input type="text" name="username" id="username" placeholder="Enter A Username"/>            
			<input type="submit" name="login_submit" value="Log In">

GleasonK

<script src="https://ajax.googleapis.com/ajax/libs/jquery/2.1.3/jquery.min.js"></script>
<script src="https://cdn.pubnub.com/pubnub-3.7.14.min.js"></script>
<script src="https://cdn.pubnub.com/webrtc/webrtc.js"></script>

GleasonK

<div id="tic-tac-box" style="float: left; width: 47%;"></div>
<div style="float: left; width: 50%;">
    <div id="video-chat" hidden="true">
		<div id="vid-box"></div>
		<button onclick="end()">End Call</button>
    </div>
    <form name="loginForm" id="login" action="#" onsubmit="return login(this);">
        	<input type="text" name="username" id="username" placeholder="Enter A Username"/>            
			<input type="submit" name="login_submit" value="Log In">
    </form>

GleasonK

<uses-feature android:name="android.hardware.camera" />
<uses-feature android:name="android.hardware.camera.autofocus" />
<uses-feature
    android:glEsVersion="0x00020000"
    android:required="true" />
<uses-permission android:name="android.permission.CAMERA" />
<uses-permission android:name="android.permission.RECORD_AUDIO" />
<uses-permission android:name="android.permission.INTERNET" />
<uses-permission android:name="android.permission.ACCESS_NETWORK_STATE" />
<uses-permission android:name="android.permission.MODIFY_AUDIO_SETTINGS" />

GleasonK

<div id="vid-box"><!-- Stream goes here --></div>

<form name="streamForm" id="stream" action="#" onsubmit="return stream(this);">
    <input type="text" name="streamname" id="streamname" placeholder="Pick a stream name!" />
    <input type="submit" name="stream_submit" value="Stream"> 
    <div id="stream-info">Watching: <span id="here-now">0</span></div>
</form>

<form name="watchForm" id="watch" action="#" onsubmit="return watch(this);">
	<input type="text" name="number" placeholder="Enter stream to join!" />