"""
With allow_soft_placement this code will work even if the assigned
GPU is not gpu:0 or even if we run on a node without a GPU.
"""

import os

import tensorflow as tf

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"


def get_n_cores():
    """Gets the assiged number of cores for this job. This is stored in
    the NSLOTS variable, If NSLOTS is not defined throw an exception.
    """
    nslots = os.getenv("NSLOTS")
    if nslots is not None:
        return int(nslots)
    raise ValueError("Environment variable NSLOTS is not defined.")


# --------------- Now start the Tensorflow code... ----------------------------
with tf.device("/gpu:0"):
    a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name="a")
    b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name="b")

# If an op is not assigned to a device then Tensorflow will pick one, which is
# typically the GPU if one is available
c = tf.matmul(a, b)

# Create the configuration for the Session
session_conf = tf.compat.v1.ConfigProto(
    intra_op_parallelism_threads=1,
    inter_op_parallelism_threads=get_n_cores(),
    allow_soft_placement=True,
    log_device_placement=True,
)

sess = tf.compat.v1.Session(config=session_conf)

# Runs the op.
print(sess.run(c))