Ikeda Map in TensorFlow

import tensorflow as tf

u = 0.9
batch_size=10**6

x, y = tf.random.uniform((batch_size, 1), 0, 0.5), tf.random.uniform((batch_size, 1), 0, 0.5)
N = 50  # Number of iterations

# Use tf.while_loop to perform iterations inside the TensorFlow graph
@tf.function()
def condition(i, x, y):
    return i < N

@tf.function()
def body(i, x, y):
    t = 0.4 - (6 / (1 + x ** 2 + y ** 2))
    x_new = 1 + u * (x * tf.cos(t) - y * tf.sin(t))
    y_new = u * (x * tf.sin(t) + y * tf.cos(t))
    return i + 1, x_new, y_new

i = tf.constant(0)
_, x_final, y_final = tf.while_loop(
    cond=condition,
    body=body,
    loop_vars=[i, x, y],
    parallel_iterations=4  # Adjust based on your system
)

WARNING:root:Limited tf.compat.v2.summary API due to missing TensorBoard installation.
WARNING:root:Limited tf.compat.v2.summary API due to missing TensorBoard installation.
WARNING:root:Limited tf.compat.v2.summary API due to missing TensorBoard installation.
WARNING:root:Limited tf.summary API due to missing TensorBoard installation.
WARNING:root:Limited tf.compat.v2.summary API due to missing TensorBoard installation.
WARNING:root:Limited tf.compat.v2.summary API due to missing TensorBoard installation.
WARNING:root:Limited tf.compat.v2.summary API due to missing TensorBoard installation.

93.2 ms ± 1.36 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
body(_, x_final, y_final)

10.5 ms ± 384 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)