def get_input_dims(grid, y_args, grid_in_dims):
    """
    Returns input dimensions

    Args:
        grid np.array((n_samples, n_grid)) or 
            (np.array(n_xgrid), np.array(n_tgrid))
            np.array((n_samples, n_tgrid, n_x1grid, n_x2grid, ..., dim_grid)): grid for 1D, 2D, or ND input, respectively
        y_args np.array((n_samples, n_grid, dim_y_args)) or 
            (np.array(n_samples, n_tgrid, n_xgrid, dim_y_args)) or
            np.array((n_samples, n_tgrid, n_x1grid, n_x2grid, ..., dim_y_args)): Additional inputs, e.g., forcing terms 
        grid_in_dims tuple(int): Indices of grid dimensions that are used as input, e.g., 0, 1 for tgrid and x1grid. -1 for all.
    Returns:
        n_grid tuple(int): Tuple of all grid lengths (n_tgrid, n_x1grid, n_x2grid, ...)
        n_tgrid int:
        n_xgrid int:
        dim_grid int:
        dim_in int:
    """
    # Calculate grid dimension for different shapes of input grids
    if len(grid) > 2: # n_samples, n_tgrid, n_x1 ..., dim_grid
        n_tgrid = grid.shape[1]
        n_xgrid = grid.shape[2]
        n_grid = grid.shape[1:-1]
    elif type(grid) is not tuple: # n_samples, n_grid
        n_tgrid = grid.shape[-1]
        n_xgrid = 0
        n_grid = (n_tgrid,)
    else:
        n_xgrid = grid[0].shape[0]
        n_tgrid = grid[1].shape[0]
        n_grid = (n_tgrid, n_xgrid)
    dim_grid = len(n_grid)

    # Calculate number of y_args inputs
    if y_args is not None:
        dim_y_args = y_args.shape[-1]
    else:
        dim_y_args = 0

    # Calculate number of grid input dimensions
    if len(grid_in_dims) == 1:
        if grid_in_dims[0] == -1:
            n_grid_in_dims = dim_grid 
        n_grid_in_dims = len(grid_in_dims)
    else:
        n_grid_in_dims = len(grid_in_dims)

    # Calculate number of total inputs
    dim_in = n_grid_in_dims + dim_y_args

    return n_grid, n_tgrid, n_xgrid, dim_grid, dim_in