# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import threading

from absl.testing import absltest
import jax
from jax import lax, numpy as np
from jax.config import config
from jax.lib import xla_client
import jax.test_util
import numpy as onp

config.parse_flags_with_absl()
FLAGS = config.FLAGS

class InfeedTest(jax.test_util.JaxTestCase):

  def testInfeed(self):
    @jax.jit
    def f(x):
      token = lax.create_token(x)
      (y,), token = lax.infeed(
          token, shape=(jax.ShapedArray((3, 4), np.float32),))
      (z,), _ = lax.infeed(
          token, shape=(jax.ShapedArray((3, 1, 1), np.float32),))
      return x + y + z

    x = onp.float32(1.5)
    y = onp.reshape(onp.arange(12, dtype=onp.float32), (3, 4)) # onp.random.randn(3, 4).astype(onp.float32)
    z = onp.random.randn(3, 1, 1).astype(onp.float32)
    xla_client.transfer_to_infeed((y,))
    xla_client.transfer_to_infeed((z,))
    self.assertAllClose(f(x), x + y + z, check_dtypes=True)

  def testInfeedThenOutfeed(self):
    @jax.jit
    def f(x):
      token = lax.create_token(x)
      y, token = lax.infeed(
          token, shape=jax.ShapedArray((3, 4), np.float32))
      token = lax.outfeed(token, y + onp.float32(1))
      return lax.tie_in(token, x - 1)

    x = onp.float32(7.5)
    y = onp.random.randn(3, 4).astype(onp.float32)
    execution = threading.Thread(target=lambda: f(x))
    execution.start()
    xla_client.transfer_to_infeed((y,))
    out, = xla_client.transfer_from_outfeed(xla_client.shape_from_pyval((y,)))
    execution.join()
    self.assertAllClose(out, y + onp.float32(1), check_dtypes=True)

  def testInfeedThenOutfeedInALoop(self):
    def doubler(_, token):
      y, token = lax.infeed(
          token, shape=jax.ShapedArray((3, 4), np.float32))
      return lax.outfeed(token, y * onp.float32(2))

    @jax.jit
    def f(n):
      token = lax.create_token(n)
      token = lax.fori_loop(0, n, doubler, token)
      return lax.tie_in(token, n)

    n = 10
    execution = threading.Thread(target=lambda: f(n))
    execution.start()
    for _ in range(n):
      x = onp.random.randn(3, 4).astype(onp.float32)
      xla_client.transfer_to_infeed((x,))
      y, = xla_client.transfer_from_outfeed(xla_client.shape_from_pyval((x,)))
      self.assertAllClose(y, x * onp.float32(2), check_dtypes=True)
    execution.join()


if __name__ == '__main__':
  absltest.main()

Add infeed test case.