ejkernel.callib._utils

ejkernel.callib._utils#

Utility functions for the callib module.

This module provides common utility functions used across the callib package, including mathematical operations, array shape calculations, system utilities, and environment configuration helpers.

Key Functions:

  • cdiv: Ceiling division for integers and JAX arrays

  • strides_from_shape: Calculate strides for contiguous arrays

  • next_power_of_2: Find next power of 2

  • get_cache_dir: Get platform-specific cache directory

  • quiet: Context manager for suppressing output

  • check_bool_flag: Parse boolean environment variables

Protocols:

  • ShapeDtype: Protocol for array-like objects with shape and dtype

class ejkernel.callib._utils.DummyStream[source]#

Bases: object

A null device-like stream that discards all writes.

Used for suppressing output by replacing stdout/stderr. All write and flush operations are no-ops.

flush(*args, **kwargs)[source]#

Discard all flush operations.

write(*args, **kwargs)[source]#

Discard all write operations.

class ejkernel.callib._utils.ShapeDtype(*args, **kwargs)[source]#

Bases: Protocol

Protocol for objects that have shape and dtype attributes.

This protocol defines the interface for array-like objects that provide shape and dtype information, commonly used in tensor operations.

property dtype: dtype#
property shape: tuple[int, ...]#
ejkernel.callib._utils.cdiv(a: int, b: int) → int[source]#
ejkernel.callib._utils.cdiv(a: int, b: Array) → Array
ejkernel.callib._utils.cdiv(a: Array, b: int) → Array
ejkernel.callib._utils.cdiv(a: Array, b: Array) → Array

Ceiling division operation.

Computes the ceiling division of a by b, which is equivalent to (a + b - 1) // b.

Parameters

  • a – Dividend, can be an integer or a JAX array.

  • b – Divisor, can be an integer or a JAX array.

Returns

The ceiling division result with the same type as inputs.

ejkernel.callib._utils.check_bool_flag(name: str, default: bool = True) → bool[source]#

Parse boolean environment variable.

Interprets various string representations as boolean values. Accepts: ‘true’, ‘yes’, ‘ok’, ‘1’, ‘easy’ (case-insensitive).

Parameters

  • name – Environment variable name.

  • default – Default value if variable not set.

Returns

Boolean interpretation of the environment variable.

Example

>>> os.environ['DEBUG'] = 'yes'
>>> check_bool_flag('DEBUG')
True
>>> check_bool_flag('MISSING', default=False)
False
ejkernel.callib._utils.get_cache_dir() → Path[source]#

Get the EasyDeL cache directory.

Returns the platform-specific cache directory for EasyDeL. Creates the directory if it doesn’t exist.

Returns

Path to the cache directory.

Example

>>> cache_dir = get_cache_dir()
>>> print(cache_dir)
/home/user/.cache/ejkernel-cache
ejkernel.callib._utils.next_power_of_2(x: int) → int[source]#

Returns the next power of two greater than or equal to x.

Parameters

x – A non-negative integer.

Returns

The smallest power of 2 greater than or equal to x.

Raises

ValueError – If x is negative.

ejkernel.callib._utils.quiet(suppress_stdout=True, suppress_stderr=True)[source]#

Context manager to temporarily suppress stdout and/or stderr output.

Replaces stdout/stderr with null streams to discard all output. Restores original streams on exit.

Parameters

  • suppress_stdout – Whether to suppress stdout.

  • suppress_stderr – Whether to suppress stderr.

Yields

None

Example

>>> with quiet():
...     print("This won't be displayed")
...     noisy_function()
>>> print("This will be displayed")

Note

This will suppress ALL output to the specified streams within the context, including output from C extensions and system calls.

ejkernel.callib._utils.strides_from_shape(shape: tuple[int, ...]) → tuple[int, ...][source]#

Calculate the strides for a contiguous array with the given shape.

Computes the number of elements to skip in memory to advance by one position along each dimension, assuming row-major (C-style) layout.

Parameters

shape – A tuple of integers representing the dimensions of an array.

Returns

A tuple of integers representing the strides of a contiguous array. The stride for dimension i is the product of all dimensions after i.

Example

>>> strides_from_shape((2, 3, 4))
(12, 4, 1)
Contents